#![allow(dead_code)]

use ash::{
    Device, Entry, Instance,
    ext::debug_utils,
    khr::swapchain::{self, *},
    prelude::VkResult,
    vk::{
        self, AccessFlags, BufferMemoryBarrier, CommandBufferLevel, CommandBufferUsageFlags,
        CommandPool, DependencyFlags, Extent3D, Handle, Image, ImageAspectFlags, ImageLayout,
        Offset3D, PFN_vkCreateWin32SurfaceKHR, PipelineStageFlags, QueueFlags, Result, SharingMode,
        SurfaceKHR, SwapchainKHR, make_api_version,
    },
};
use getset::{CloneGetters, CopyGetters, Getters, MutGetters, Setters, WithSetters};
use std::{
    default,
    ffi::{c_char, CStr, CString},
    mem::*,
    os::raw::*,
    ptr::{null, null_mut},
    sync::{Arc, LazyLock, Mutex, RwLock}, u64,
};
use wgpu::rwh::{HasDisplayHandle, HasWindowHandle, RawWindowHandle};
use winit::{dpi::*, platform::windows::*, window::Window};

use crate::engine::njm_device::{NjmDevice, get_device_proc_addr, get_device_proc_addr_converted};

pub static MAX_FRAMES_IN_FLIGHT: u32 = 2;
static NJM_IMAGE_FORMAT: vk::Format = vk::Format::B8G8R8A8_SRGB;
static NJM_DEP_STENCIL_FORMAT: vk::Format = vk::Format::D32_SFLOAT_S8_UINT;
static NJM_IMAGE_COLOR_SPACE: vk::ColorSpaceKHR = vk::ColorSpaceKHR::SRGB_NONLINEAR;
static NJM_PRESENT_MODE: vk::PresentModeKHR = vk::PresentModeKHR::FIFO;

pub struct NjmDepthResource {
    dep_img: vk::Image,
    dep_mem: vk::DeviceMemory,
    dep_view: vk::ImageView,
}

pub struct NjmSyncObj {
    img_avail_sem: vk::Semaphore,
    render_finished_sem: vk::Semaphore,
    inflight_fence: vk::Fence,
}

#[derive(Getters, Setters, WithSetters, MutGetters, CopyGetters, CloneGetters)]
pub struct NjmSwapChain {
    #[getset(get = "pub")]
    device: Arc<Mutex<NjmDevice>>,
    
    #[getset(get = "pub")]
    wnd_extent: vk::Extent2D,

    #[getset(get = "pub")]
    swap_chain: vk::SwapchainKHR,

    #[getset(get = "pub")]
    img_fmt: vk::Format,

    #[getset(get = "pub")]
    dep_fmt: vk::Format,
    
    #[getset(get = "pub")]
    sc_extent: vk::Extent2D,

    #[getset(get = "pub")]
    render_pass: vk::RenderPass,

    #[getset(get = "pub")]
    frame_bufs: Vec<vk::Framebuffer>,

    #[getset(get = "pub")]
    imgs: Vec<vk::Image>,

    #[getset(get = "pub")]
    img_views: Vec<vk::ImageView>,

    // dep_img: vk::Image,
    // dep_mem: vk::DeviceMemory,
    // dep_view: vk::ImageView,
    #[getset(get = "pub")]
    dep_res: Vec<NjmDepthResource>,

    // img_avail_sems: Vec<vk::Semaphore>,
    // rend_finished_sems: Vec<vk::Semaphore>,
    // inflight_fens: Vec<vk::Fence>,
    #[getset(get = "pub")]
    sync_objs: Vec<NjmSyncObj>,

    #[getset(get = "pub")]
    imgs_inflight: Vec<vk::Fence>,

    #[getset(get = "pub", set = "pub")]
    current_frame: usize,
}

impl Drop for NjmSwapChain {
    fn drop(&mut self) {
        let dev = self.device.lock().unwrap();

        for img_view in &self.img_views {
            if !img_view.is_null() {
                unsafe { dev.device().destroy_image_view(*img_view, None) };
            }
        }
        self.img_views.clear();

        for res in &self.dep_res {
            if !res.dep_view.is_null() {
                unsafe { dev.device().destroy_image_view(res.dep_view, None) };
            }
            if !res.dep_img.is_null() {
                unsafe { dev.device().destroy_image(res.dep_img, None) };
            }
            if !res.dep_mem.is_null() {
                unsafe { dev.device().free_memory(res.dep_mem, None) };
            }
        }
        self.dep_res.clear();

        for sync_obj in &self.sync_objs {
            if !sync_obj.img_avail_sem.is_null() {
                unsafe { dev.device().destroy_semaphore(sync_obj.img_avail_sem, None) };
            }
            if !sync_obj.render_finished_sem.is_null() {
                unsafe { dev.device().destroy_semaphore(sync_obj.render_finished_sem, None) };
            }
            if !sync_obj.inflight_fence.is_null() {
                unsafe { dev.device().destroy_fence(sync_obj.inflight_fence, None) };
            }
        }
        self.sync_objs.clear();
        self.imgs_inflight.clear(); // ?

        for framebuf in &self.frame_bufs {
            if !framebuf.is_null() {
                unsafe { dev.device().destroy_framebuffer(*framebuf, None) };
            }
        }
        self.frame_bufs.clear();

        if !self.render_pass.is_null() {
            unsafe { dev.device().destroy_render_pass(self.render_pass, None) };
        }

        if !self.swap_chain.is_null() {
            unsafe { dev.get_khr_swapchain_device().destroy_swapchain(self.swap_chain, None) };
        }
    }
}

impl NjmSwapChain {

    // public api
    pub fn aquire_next_image(&self) -> (u32, bool) {
        let current_sync_obj = &self.sync_objs[self.current_frame];
        let wait_fences = [current_sync_obj.inflight_fence];
        let sem = current_sync_obj.img_avail_sem;

        let device_ref = self.device.lock().unwrap();
        let _ = unsafe { device_ref.device().wait_for_fences(&wait_fences, true, u64::MAX).unwrap() };
        let _ = unsafe { device_ref.device().reset_fences(&wait_fences).unwrap() };
        let rst = unsafe { device_ref.get_khr_swapchain_device().acquire_next_image(self.swap_chain, u64::MAX, sem, vk::Fence::null()).unwrap() };
        return rst;
    }

    pub fn submit_command_buffers(&mut self, buffers: &[vk::CommandBuffer], img_idx: usize) -> bool { // fatal problem!
        let device_ref = self.device.lock().unwrap();

        self.imgs_inflight[img_idx] = self.sync_objs[self.current_frame].inflight_fence;
        let que = device_ref.get_graphics_queue();
        let wait_sems = [self.sync_objs[self.current_frame].img_avail_sem];
        let signal_sems = [self.sync_objs[self.current_frame].render_finished_sem];
        let dst_stage_flags = [vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT];
        
        let sub_info = vk::SubmitInfo::default()
            .command_buffers(buffers)
            .wait_semaphores(&wait_sems)
            .signal_semaphores(&signal_sems)
            .wait_dst_stage_mask(&dst_stage_flags);
        let subs = [sub_info];
        let _ = unsafe { device_ref.device().queue_submit(que, &subs, self.imgs_inflight[img_idx]).unwrap() };

        let scs = [self.swap_chain];
        let img_idxs = [img_idx as u32];
        let mut rsts = [vk::Result::SUCCESS];
        let present_info = vk::PresentInfoKHR::default()
            .swapchains(&scs)
            .wait_semaphores(&signal_sems)
            .image_indices(&img_idxs)
            .results(&mut rsts);
        let r: bool = unsafe { device_ref.get_khr_swapchain_device().queue_present(que, &present_info).unwrap() };
        if rsts[0] != vk::Result::SUCCESS {
            println!("swapchain::submit_command_buffers::queue_present : present results is {:?}", rsts);
        }
        return r;
    }

    // init
    pub fn new(
        device: Arc<Mutex<NjmDevice>>,
        wnd_extent: vk::Extent2D,
        previous: Option<Arc<Mutex<NjmSwapChain>>>,
    ) -> Self {
        let dev_ref = device.lock().unwrap();

        let swap_chain = Self::create_swap_chain(&dev_ref, wnd_extent, previous);
        let imgs = Self::get_swapchain_images(&dev_ref, &swap_chain);
        let img_cnt = imgs.len();
        let views = Self::create_image_views(&dev_ref, &imgs);
        let render_pass = Self::create_render_pass(&dev_ref);
        let dep_res = Self::create_depth_resources(&dev_ref, &wnd_extent, img_cnt);
        let fb_bufs = Self::create_frame_buffers(&dev_ref, &render_pass, &wnd_extent, &views, &dep_res);
        let sync_objs = Self::create_sync_objs(&dev_ref);
        let mut imgs_inflight = vec![];
        let null_fence = vk::Fence::null();
        for _i in 0..img_cnt {
            imgs_inflight.push(null_fence);
        }
        
        return Self {
            device: device.clone(),
            wnd_extent: wnd_extent,
            swap_chain: swap_chain,
            img_fmt: NJM_IMAGE_FORMAT,
            dep_fmt: NJM_DEP_STENCIL_FORMAT,
            sc_extent: wnd_extent,
            imgs: imgs,
            img_views: views,
            render_pass: render_pass,
            dep_res: dep_res,
            frame_bufs: fb_bufs,
            sync_objs: sync_objs,
            imgs_inflight: imgs_inflight,
            current_frame: 0,
        };
    }

    // init helper
    fn create_swap_chain(
        device: &NjmDevice,
        wnd_extent: vk::Extent2D,
        previous: Option<Arc<Mutex<NjmSwapChain>>>,
    ) -> vk::SwapchainKHR {
        let old_sc: vk::SwapchainKHR = if let Some(sc) = previous {
            sc.lock().unwrap().swap_chain
        } else {
            vk::SwapchainKHR::null()
        };
        let sc_info = vk::SwapchainCreateInfoKHR::default()
            .surface(*device.surface())
            .min_image_count(2)
            .image_format(NJM_IMAGE_FORMAT)
            .image_color_space(NJM_IMAGE_COLOR_SPACE)
            .image_extent(wnd_extent)
            .image_array_layers(1)
            .image_usage(vk::ImageUsageFlags::COLOR_ATTACHMENT)
            .pre_transform(vk::SurfaceTransformFlagsKHR::IDENTITY)
            .composite_alpha(vk::CompositeAlphaFlagsKHR::OPAQUE)
            .present_mode(NJM_PRESENT_MODE)
            .clipped(true)
            .old_swapchain(old_sc);
        let sc = unsafe {
            device
                .get_khr_swapchain_device()
                .create_swapchain(&sc_info, None)
                .unwrap()
        };
        return sc;
    }

    fn get_swapchain_images(device: &NjmDevice, swap_chain: &SwapchainKHR) -> Vec<vk::Image> {
        let dev = device.get_khr_swapchain_device();
        let imgs = unsafe { dev.get_swapchain_images(*swap_chain).unwrap() };
        return imgs;
    }

    fn create_image_views(device: &NjmDevice, imgs: &Vec<vk::Image>) -> Vec<vk::ImageView> {
        let mut views: Vec<vk::ImageView> = vec![];
        for img in imgs {
            let comp_swzl = vk::ComponentSwizzle::IDENTITY;
            let comp_map = vk::ComponentMapping::default()
                .r(comp_swzl)
                .g(comp_swzl)
                .b(comp_swzl)
                .a(comp_swzl);
            let sub_res_range = vk::ImageSubresourceRange::default()
                .aspect_mask(ImageAspectFlags::COLOR)
                .base_array_layer(0)
                .base_mip_level(0)
                .layer_count(1)
                .level_count(1);
            let view_info = vk::ImageViewCreateInfo::default()
                .format(NJM_IMAGE_FORMAT)
                .view_type(vk::ImageViewType::TYPE_2D)
                .image(*img)
                .components(comp_map)
                .subresource_range(sub_res_range);

            let view = unsafe { device.device().create_image_view(&view_info, None).unwrap() };
            views.push(view);
        }
        return views;
    }

    fn create_render_pass(device: &NjmDevice) -> vk::RenderPass {
        let color_attach_desc = vk::AttachmentDescription::default()
            .format(NJM_IMAGE_FORMAT)
            .samples(vk::SampleCountFlags::TYPE_1)
            .load_op(vk::AttachmentLoadOp::CLEAR)
            .store_op(vk::AttachmentStoreOp::STORE)
            .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
            .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
            .initial_layout(vk::ImageLayout::UNDEFINED)
            .final_layout(vk::ImageLayout::PRESENT_SRC_KHR);
        let depstencil_attach_desc = vk::AttachmentDescription::default()
            .format(NJM_DEP_STENCIL_FORMAT)
            .samples(vk::SampleCountFlags::TYPE_1)
            .load_op(vk::AttachmentLoadOp::CLEAR)
            .store_op(vk::AttachmentStoreOp::DONT_CARE)
            .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
            .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
            .initial_layout(vk::ImageLayout::UNDEFINED)
            .final_layout(vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL);
        let attachs = [color_attach_desc, depstencil_attach_desc];

        let color_attach_ref = vk::AttachmentReference::default()
            .attachment(0)
            .layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL);
        let depstencil_attach_ref = vk::AttachmentReference::default()
            .attachment(1)
            .layout(vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL);
        let color_attach_refs = [color_attach_ref];

        let subpass = vk::SubpassDescription::default()
            .color_attachments(&color_attach_refs)
            .depth_stencil_attachment(&depstencil_attach_ref)
            .pipeline_bind_point(vk::PipelineBindPoint::GRAPHICS);
        let subpass_dep = vk::SubpassDependency::default()
            .src_subpass(vk::SUBPASS_EXTERNAL)
            .src_access_mask(vk::AccessFlags::NONE)
            .src_stage_mask(vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT)
            .dst_subpass(0)
            .dst_access_mask(vk::AccessFlags::COLOR_ATTACHMENT_WRITE)
            .dst_stage_mask(vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT);
        let subpasses = [subpass];
        let subpass_deps = [subpass_dep];

        let render_pass_info = vk::RenderPassCreateInfo::default()
            .attachments(&attachs)
            .subpasses(&subpasses)
            .dependencies(&subpass_deps);
        let render_pass = unsafe {
            device
                .device()
                .create_render_pass(&render_pass_info, None)
                .unwrap()
        };
        return render_pass;
    }

    fn create_depth_resources(
        device: &NjmDevice,
        extent: &vk::Extent2D,
        count: usize,
    ) -> Vec<NjmDepthResource> {
        let mut dep_res: Vec<NjmDepthResource> = vec![];
        for i in 0..count {
            let depimg_info = vk::ImageCreateInfo::default()
                .image_type(vk::ImageType::TYPE_2D)
                .extent(vk::Extent3D {
                    width: extent.width,
                    height: extent.height,
                    depth: 1,
                })
                .mip_levels(1)
                .array_layers(1)
                .format(NJM_DEP_STENCIL_FORMAT)
                .tiling(vk::ImageTiling::OPTIMAL)
                .initial_layout(vk::ImageLayout::UNDEFINED)
                .usage(vk::ImageUsageFlags::DEPTH_STENCIL_ATTACHMENT)
                .samples(vk::SampleCountFlags::TYPE_1)
                .sharing_mode(vk::SharingMode::EXCLUSIVE)
                .flags(vk::ImageCreateFlags::empty());
            let (img, mem) = device
                .create_image(depimg_info, vk::MemoryPropertyFlags::DEVICE_LOCAL)
                .unwrap();

            let comp_swzl = vk::ComponentSwizzle::IDENTITY;
            let comp_map = vk::ComponentMapping::default()
                .r(comp_swzl)
                .g(comp_swzl)
                .b(comp_swzl)
                .a(comp_swzl);
            let sub_res_range = vk::ImageSubresourceRange::default()
                .aspect_mask(ImageAspectFlags::DEPTH)
                .base_array_layer(0)
                .base_mip_level(0)
                .layer_count(1)
                .level_count(1);
            let view_info = vk::ImageViewCreateInfo::default()
                .format(NJM_DEP_STENCIL_FORMAT)
                .view_type(vk::ImageViewType::TYPE_2D)
                .image(img)
                .components(comp_map)
                .subresource_range(sub_res_range);
            let view = unsafe { device.device().create_image_view(&view_info, None).unwrap() };

            dep_res.push(NjmDepthResource {
                dep_img: img,
                dep_mem: mem,
                dep_view: view,
            });
        }
        return dep_res;
    }

    fn create_frame_buffers(
        device: &NjmDevice,
        render_pass: &vk::RenderPass,
        extent: &vk::Extent2D,
        img_views: &Vec<vk::ImageView>,
        dep_res: &Vec<NjmDepthResource>,
    ) -> Vec<vk::Framebuffer> {
        assert_eq!(img_views.len(), dep_res.len());
        let cnt = img_views.len();

        let mut frame_bufs: Vec<vk::Framebuffer> = vec![];
        for i in 0..cnt {
            let attachs = [img_views[i], dep_res[i].dep_view];
            let fb_info = vk::FramebufferCreateInfo::default()
                .render_pass(*render_pass)
                .attachment_count(attachs.len() as u32)
                .attachments(&attachs)
                .width(extent.width)
                .height(extent.height)
                .layers(1);
            let fb = unsafe { device.device().create_framebuffer(&fb_info, None).unwrap() };
            frame_bufs.push(fb);
        }
        return frame_bufs;
    }

    fn create_sync_objs(device: &NjmDevice) -> Vec<NjmSyncObj> {
        let sem_info = vk::SemaphoreCreateInfo::default()
            .flags(vk::SemaphoreCreateFlags::empty());
        let fence_info = vk::FenceCreateInfo::default()
            .flags(vk::FenceCreateFlags::SIGNALED);

        let mut sync_objs: Vec<NjmSyncObj> = vec![];
        for i in 0..MAX_FRAMES_IN_FLIGHT {
            let img_avail_sem = unsafe { device.device().create_semaphore(&sem_info, None).unwrap() };
            let render_finished_sem = unsafe { device.device().create_semaphore(&sem_info, None).unwrap() };
            let inflight_fence = unsafe { device.device().create_fence(&fence_info, None).unwrap() };
            let sync_obj = NjmSyncObj {
                img_avail_sem: img_avail_sem,
                render_finished_sem: render_finished_sem,
                inflight_fence: inflight_fence,
            };
            sync_objs.push(sync_obj);
        }
        return sync_objs;
    }

    // choose helper
    pub fn choose_swapchain_present_mode(availables: Vec<vk::PresentModeKHR>) -> vk::PresentModeKHR {
        for mode in availables {
            if mode == vk::PresentModeKHR::MAILBOX {
                return mode;
            }
        }
        return vk::PresentModeKHR::FIFO;
    }

    pub fn choose_swapchain_surface_format(availables: Vec<vk::SurfaceFormatKHR>) -> vk::SurfaceFormatKHR {
        for format in &availables {
            if format.format == vk::Format::B8G8R8A8_SRGB &&
                format.color_space == vk::ColorSpaceKHR::SRGB_NONLINEAR 
            {
                return *format;
            }
        }
        return availables[0];
    }

    pub fn choose_swapchain_extent(wnd_extent: vk::Extent2D, caps: vk::SurfaceCapabilitiesKHR) -> vk::Extent2D {
        if caps.current_extent.width != u32::MAX {
            return caps.current_extent;
        } else {
            let actual_extent: vk::Extent2D = vk::Extent2D { 
                width: wnd_extent.width.clamp(caps.min_image_extent.width, caps.max_image_extent.width), 
                height: wnd_extent.height.clamp(caps.min_image_extent.height, caps.max_image_extent.height),
            };
            return actual_extent;
        }
    }

    pub fn choose_supported_format(
        device: &NjmDevice,
        availables: Vec<vk::Format>,
        tiling: vk::ImageTiling,
        features: vk::FormatFeatureFlags,
    ) -> vk::Format {
        for format in &availables {
            let property = unsafe { device.instance().get_physical_device_format_properties(*device.physical_device(), *format) };
            if tiling == vk::ImageTiling::LINEAR && 
                (property.linear_tiling_features & features) == features {
                return *format;
            } else if tiling == vk::ImageTiling::OPTIMAL &&
                (property.optimal_tiling_features & features) == features {
                return *format;
            }
        }
        panic!("unable to find supported format!");
    }

    pub fn choose_depth_format(device: &NjmDevice) -> vk::Format {
        let availables = vec![
            vk::Format::D32_SFLOAT,
            vk::Format::D32_SFLOAT_S8_UINT,
            vk::Format::D24_UNORM_S8_UINT,
        ];
        let tiling = vk::ImageTiling::OPTIMAL;
        let features = vk::FormatFeatureFlags::DEPTH_STENCIL_ATTACHMENT;
        return Self::choose_supported_format(device, availables, tiling, features);
    }

    // other helper
    pub fn compare_swap_formats(&self, other: &NjmSwapChain) -> bool {
        return other.dep_fmt == self.dep_fmt && other.img_fmt == self.img_fmt;
    }

    pub fn get_swapchain_aspect_ratio(&self) -> f32 {
        return self.sc_extent.width as f32 / self.sc_extent.height as f32;
    }
}
